import os
import pandas
import numpy
import math
from torch.utils.data import Dataset

def load_UCR_dataset(root_path,
                     data_path,
                     data_name):

    ucrlist = os.listdir(os.path.join(root_path, data_path))
    if data_name not in ucrlist:
        raise ValueError('dataset not found')
    train_file = os.path.join(root_path, data_path, data_name, data_name + "_TRAIN.tsv")
    test_file = os.path.join(root_path, data_path, data_name, data_name + "_TEST.tsv")
    train_df = pandas.read_csv(train_file, sep='\t', header=None)
    test_df = pandas.read_csv(test_file, sep='\t', header=None)
    train_array = numpy.array(train_df)
    test_array = numpy.array(test_df)

    # Move the labels to {0, ..., L-1}
    labels = numpy.unique(train_array[:, 0])
    transform = {}
    for i, l in enumerate(labels):
        transform[l] = i

    train = numpy.expand_dims(train_array[:, 1:], 1).astype(numpy.float64)
    train_labels = numpy.vectorize(transform.get)(train_array[:, 0])
    test = numpy.expand_dims(test_array[:, 1:], 1).astype(numpy.float64)
    test_labels = numpy.vectorize(transform.get)(test_array[:, 0])

    # Normalization for non-normalized datasets
    # To keep the amplitude information, we do not normalize values over
    # individual time series, but on the whole dataset
    if data_name not in [
        'AllGestureWiimoteX',
        'AllGestureWiimoteY',
        'AllGestureWiimoteZ',
        'BME',
        'Chinatown',
        'Crop',
        'EOGHorizontalSignal',
        'EOGVerticalSignal',
        'Fungi',
        'GestureMidAirD1',
        'GestureMidAirD2',
        'GestureMidAirD3',
        'GesturePebbleZ1',
        'GesturePebbleZ2',
        'GunPointAgeSpan',
        'GunPointMaleVersusFemale',
        'GunPointOldVersusYoung',
        'HouseTwenty',
        'InsectEPGRegularTrain',
        'InsectEPGSmallTrain',
        'MelbournePedestrian',
        'PickupGestureWiimoteZ',
        'PigAirwayPressure',
        'PigArtPressure',
        'PigCVP',
        'PLAID',
        'PowerCons',
        'Rock',
        'SemgHandGenderCh2',
        'SemgHandMovementCh2',
        'SemgHandSubjectCh2',
        'ShakeGestureWiimoteZ',
        'SmoothSubspace',
        'UMD'
    ]:
        return train, train_labels, test, test_labels
    # Post-publication note:
    # Using the testing set to normalize might bias the learned network,
    # but with a limited impact on the reported results on few datasets.
    # See the related discussion here: https://github.com/White-Link/UnsupervisedScalableRepresentationLearningTimeSeries/pull/13.
    mean = numpy.nanmean(numpy.concatenate([train, test]))
    var = numpy.nanvar(numpy.concatenate([train, test]))
    train = (train - mean) / math.sqrt(var)
    test = (test - mean) / math.sqrt(var)
    return train, train_labels, test, test_labels

class Dataset_UCR(Dataset):
    def __init__(self,
                 root_path,
                 data_path,
                 data_name,
                 flag='TRAIN',
                 config=None
                ):

        # init
        self.root_path = root_path
        self.data_path = data_path
        self.data_name = data_name
        self.flag = flag

        self.config = config

        self.__read_data()


    def __read_data(self):
        if self.flag == 'TRAIN':
            self.x, self.y, *self._ = load_UCR_dataset(
                root_path=self.root_path,
                data_path=self.data_path,
                data_name=self.data_name)
        elif self.flag == 'TEST':
            *self._, self.x, self.y,  = load_UCR_dataset(
                root_path=self.root_path,
                data_path=self.data_path,
                data_name=self.data_name)
        else:
            raise ValueError('Unknown flag, TRAIN or TEST is required')

    def __getitem__(self, idx):
        seq_x = self.x[idx].T
        seq_y = self.y[idx]
        return seq_x, seq_y

    def __len__(self):
        return self.x.shape[0]


if __name__ == "__main__":
    '''
    ['Phoneme', 'FaceAll', 'PhalangesOutlinesCorrect', 'Crop', 'ProximalPhalanxOutlineAgeGroup', 'TwoLeadECG', 'Computers', 'MixedShapesSmallTrain', 
    'ShakeGestureWiimoteZ', 'UWaveGestureLibraryY', 'Earthquakes', 'PowerCons', 'MiddlePhalanxTW', 'ProximalPhalanxTW', 'Meat', 'FreezerRegularTrain', 
    'EthanolLevel', 'PigArtPressure', 'NonInvasiveFetalECGThorax2', 'GestureMidAirD1', 'CricketX', 'Chinatown', 'Adiac', 'BME', 'DistalPhalanxOutlineAgeGroup', 
    'SemgHandSubjectCh2', 'WormsTwoClass', 'GesturePebbleZ1', 'CinCECGTorso', 'Worms', 'ToeSegmentation2', 'GestureMidAirD2', 'SmoothSubspace', 'InlineSkate', 
    'ToeSegmentation1', 'SemgHandMovementCh2', 'MiddlePhalanxOutlineCorrect', 'TwoPatterns', 'SonyAIBORobotSurface1', 'EOGHorizontalSignal', 'GestureMidAirD3', 
    'Trace', 'InsectEPGSmallTrain', 'BeetleFly', 'ProximalPhalanxOutlineCorrect', 'OSULeaf', 'AllGestureWiimoteX', 'DodgerLoopGame', 'GunPointOldVersusYoung', 
    'DistalPhalanxOutlineCorrect', 'DistalPhalanxTW', 'StarLightCurves', 'Strawberry', 'Haptics', 'AllGestureWiimoteY', 'FreezerSmallTrain', 'Rock', 
    'WordSynonyms', 'ECG5000', 'Mallat', 'SwedishLeaf', 'MiddlePhalanxOutlineAgeGroup', 'RefrigerationDevices', 'InsectWingbeatSound', 'PickupGestureWiimoteZ', 
    'Beef', 'MoteStrain', 'MedicalImages', 'ItalyPowerDemand', 'Herring', 'SmallKitchenAppliances', 'Lightning2', 'FaceFour', 'FordB', 'CBF', 'GunPoint', 
    'UWaveGestureLibraryX', 'FiftyWords', 'ChlorineConcentration', 'DiatomSizeReduction', 'BirdChicken', 'ElectricDevices', 'DodgerLoopWeekend', 'MelbournePedestrian', 
    'PigCVP', 'ScreenType', 'PigAirwayPressure', 'UWaveGestureLibraryZ', 'ECG200', 'UWaveGestureLibraryAll', 'HandOutlines', 'Wafer', 'FacesUCR', 'Car', 'OliveOil', 
    'Yoga', 'ShapeletSim', 'HouseTwenty', 'SyntheticControl', 'UMD', 'SonyAIBORobotSurface2', 'FordA', 'CricketY', 'Plane', 'DodgerLoopDay', 'ArrowHead', 'Ham', 
    'Symbols', 'PLAID', 'CricketZ', 'ACSF1', 'GunPointMaleVersusFemale', 'ShapesAll', 'Wine', 'ECGFiveDays', 'NonInvasiveFetalECGThorax1', 'Coffee', 'Lightning7', 
    'GesturePebbleZ2', 'MixedShapesRegularTrain', 'Fungi', 'GunPointAgeSpan', 'LargeKitchenAppliances', 'InsectEPGRegularTrain', 'AllGestureWiimoteZ', 
    'EOGVerticalSignal', 'SemgHandGenderCh2', 'Fish']
    '''
    root_path = '../data'
    data_path = 'UCR/UCRArchive_2018'
    data_name = 'Beef'

    train, train_labels, test, test_labels = load_UCR_dataset(
        root_path=root_path,
        data_path=data_path,
        data_name=data_name)
    print('-[train shape]:', train.shape)
    print('-[train_labels shape]:', train_labels.shape)
    print('-[test shape]:', test.shape)
    print('-[test_labels shape]:', test_labels.shape)

    ucrtrain = Dataset_UCR(
        root_path=root_path,
        data_path=data_path,
        data_name=data_name,
        flag='TRAIN')
    print('-[UCR data class]:', ucrtrain)
    print('-[Num of data]', len(ucrtrain))
    x, y = ucrtrain[0]
    print('-[single data shape]:', x.shape)
    print('-[dim of data]:', x.shape[1])
    print('-[length of data]:', x.shape[0])
    print('-[single data label]:', y)
